#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# 导入PIL库中的Image模块，用于处理图像
import PIL.Image
# 导入matplotlib库中的pyplot模块，用于绘制和显示图像
import matplotlib.pyplot as plt
# 从torchvision.datasets中导入VOCSegmentation类，用于加载VOC数据集进行语义分割任务
from torchvision.datasets import VOCSegmentation
# 导入numpy库，用于处理数组和矩阵
import numpy as np

# 从本地文件下载或加载用于语义分割任务的VOC2012训练集
voc2007_seg_trainval = VOCSegmentation(
    root='./data',  # 数据集存储的根目录
    year='2012',  # 指定数据集的年份为2012
    image_set='train',  # 使用训练集
    download=False,  # 不进行数据集下载，如果已经下载好则设置为False
)

# 定义一个函数，用于在图像上绘制分割掩码
def draw_mask(image: PIL.Image, target: PIL.Image = None) -> PIL.Image:
    """在PASCAL VOC图像上绘制分割掩码"""
    # 定义一个颜色映射字典，用于为不同的分割类别分配颜色
    color_map = {
        0: (0, 0, 0, 128),  # 背景，黑色
        1: (247, 116, 95, 128),  # 飞机
        2: (232, 129, 49, 128),  # 自行车
        3: (208, 142, 49, 128),  # 鸟
        4: (190, 150, 49, 128),  # 船
        5: (173, 156, 49, 128),  # 瓶子
        6: (173, 156, 49, 128),  # 公交车
        7: (155, 162, 49, 128),  # 汽车
        8: (134, 167, 49, 128),  # 猫
        9: (99, 174, 49, 128),  # 椅子
        10: (49, 178, 82, 128),  # 牛
        11: (51, 176, 122, 128),  # 餐桌
        12: (52, 174, 142, 128),  # 狗
        13: (53, 173, 157, 128),  # 马
        14: (54, 172, 170, 128),  # 摩托车
        15: (54, 170, 182, 128),  # 人
        16: (56, 168, 197, 128),  # 盆栽植物
        17: (57, 166, 216, 128),  # 羊
        18: (73, 160, 244, 128),  # 沙发
        19: (135, 149, 244, 128),  # 火车
        20: (172, 136, 244, 128),  # 电视/显示器
        255: (255, 255, 255, 128),  # 未标记区域，白色
    }
    # 创建一个新的RGBA模式的掩码图像，初始全透明
    mask = PIL.Image.new("RGBA", image.size, (0, 0, 0, 0))

    # 将目标图像（掩码）转换为numpy数组，这里将数组进行了转置
    # 疑问：为什么将PIL.PngImageFile转换为numpy数组时，宽度和高度会颠倒？
    target_array = np.array(target).T

    # 遍历目标图像的每个像素
    for x in range(target.width):
        for y in range(target.height):
            # 获取当前像素的掩码值
            mpv = target_array[x, y]
            # 如果当前像素不是背景（值不为0）且不是未标记区域（值不为255）
            if target_array[x, y] != 0 and target_array[x, y] != 255:
                # 根据颜色映射字典为该像素设置颜色，透明度为128表示半透明
                mask.putpixel((x, y), color_map[mpv])

    # 将原始图像转换为RGBA模式
    image = image.convert("RGBA")
    # 将掩码图像与原始图像进行合并
    merged_image = PIL.Image.alpha_composite(image, mask)

    return merged_image

# 显示训练集中的前四张图像及其掩码
if __name__ == "__main__":
    # 创建一个大小为8x6的图像窗口
    fig = plt.figure(figsize=(8, 6))
    # 设置图像窗口的标题
    fig.suptitle("The first four images and masks"
                " \nof VOC2007 train dataset for Segmentation in PyTorch")
    # 遍历VOC2012训练集
    for idx, (image, target) in enumerate(voc2007_seg_trainval):
        # 在2x2的子图布局中添加一个子图
        ax = fig.add_subplot(2, 2, idx + 1)
        # 调用draw_mask函数绘制掩码并得到合并后的图像
        image_show = draw_mask(image, target)
        # 在子图中显示合并后的图像
        ax.imshow(image_show)

        # 禁用子图的坐标轴刻度和边框
        ax.set_frame_on(False)
        ax.set_xticks([])
        ax.set_yticks([])
        # 当显示了四张图像后停止循环
        if idx >= 3:
            break

    # 保存绘制好的图像，去除多余的空白边距
    # plt.savefig("PASCAL VOC SegementationClass.jpg", bbox_inches='tight')
    # 取消注释下面这行代码可以直接显示图像
    # plt.show()